"""
电磁场物理方程模块
"""

import torch
import torch.nn.functional as F
import numpy as np


class MaxwellEquations:
    """麦克斯韦方程组实现"""
    
    def __init__(self, eps0=1.0, mu0=1.0, sigma=0.0):
        """初始化麦克斯韦方程组参数
        
        参数:
            eps0: 真空介电常数
            mu0: 真空磁导率
            sigma: 电导率
        """
        self.eps0 = eps0
        self.mu0 = mu0
        self.sigma = sigma
    
    def compute_curl_E(self, E, dx=1.0, dy=1.0, dz=1.0):
        """计算电场旋度 ∇×E
        
        参数:
            E: 电场张量，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            电场旋度，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
        """
        # E = (Ex, Ey, Ez)
        Ex = E[..., 0]
        Ey = E[..., 1]
        Ez = E[..., 2]
        
        # 计算偏导数
        dEz_dy = torch.gradient(Ez, dim=2)[0] / dy
        dEy_dz = torch.gradient(Ey, dim=3)[0] / dz
        dEx_dz = torch.gradient(Ex, dim=3)[0] / dz
        dEz_dx = torch.gradient(Ez, dim=1)[0] / dx
        dEy_dx = torch.gradient(Ey, dim=1)[0] / dx
        dEx_dy = torch.gradient(Ex, dim=2)[0] / dy
        
        # 旋度分量
        curl_E_x = dEz_dy - dEy_dz
        curl_E_y = dEx_dz - dEz_dx
        curl_E_z = dEy_dx - dEx_dy
        
        return torch.stack([curl_E_x, curl_E_y, curl_E_z], dim=-1)
    
    def compute_curl_B(self, B, dx=1.0, dy=1.0, dz=1.0):
        """计算磁场旋度 ∇×B
        
        参数:
            B: 磁场张量，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            磁场旋度，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
        """
        # B = (Bx, By, Bz)
        Bx = B[..., 0]
        By = B[..., 1]
        Bz = B[..., 2]
        
        # 计算偏导数
        dBz_dy = torch.gradient(Bz, dim=2)[0] / dy
        dBy_dz = torch.gradient(By, dim=3)[0] / dz
        dBx_dz = torch.gradient(Bx, dim=3)[0] / dz
        dBz_dx = torch.gradient(Bz, dim=1)[0] / dx
        dBy_dx = torch.gradient(By, dim=1)[0] / dx
        dBx_dy = torch.gradient(Bx, dim=2)[0] / dy
        
        # 旋度分量
        curl_B_x = dBz_dy - dBy_dz
        curl_B_y = dBx_dz - dBz_dx
        curl_B_z = dBy_dx - dBx_dy
        
        return torch.stack([curl_B_x, curl_B_y, curl_B_z], dim=-1)
    
    def compute_div_E(self, E, dx=1.0, dy=1.0, dz=1.0):
        """计算电场散度 ∇·E
        
        参数:
            E: 电场张量，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            电场散度，形状为 [batch_size, x_dim, y_dim, z_dim]
        """
        # E = (Ex, Ey, Ez)
        Ex = E[..., 0]
        Ey = E[..., 1]
        Ez = E[..., 2]
        
        # 计算偏导数
        dEx_dx = torch.gradient(Ex, dim=1)[0] / dx
        dEy_dy = torch.gradient(Ey, dim=2)[0] / dy
        dEz_dz = torch.gradient(Ez, dim=3)[0] / dz
        
        # 散度
        div_E = dEx_dx + dEy_dy + dEz_dz
        
        return div_E
    
    def compute_div_B(self, B, dx=1.0, dy=1.0, dz=1.0):
        """计算磁场散度 ∇·B
        
        参数:
            B: 磁场张量，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            磁场散度，形状为 [batch_size, x_dim, y_dim, z_dim]
        """
        # B = (Bx, By, Bz)
        Bx = B[..., 0]
        By = B[..., 1]
        Bz = B[..., 2]
        
        # 计算偏导数
        dBx_dx = torch.gradient(Bx, dim=1)[0] / dx
        dBy_dy = torch.gradient(By, dim=2)[0] / dy
        dBz_dz = torch.gradient(Bz, dim=3)[0] / dz
        
        # 散度
        div_B = dBx_dx + dBy_dy + dBz_dz
        
        return div_B
    
    def faraday_law(self, E, B, dt=1.0, dx=1.0, dy=1.0, dz=1.0):
        """法拉第电磁感应定律: ∂B/∂t = -∇×E
        
        参数:
            E: 电场张量
            B: 磁场张量
            dt: 时间步长
            dx, dy, dz: 空间步长
            
        返回:
            磁场时间导数
        """
        curl_E = self.compute_curl_E(E, dx, dy, dz)
        dB_dt = -curl_E
        return dB_dt
    
    def ampere_law(self, E, B, J=None, dt=1.0, dx=1.0, dy=1.0, dz=1.0):
        """安培环路定律: ∂E/∂t = c²∇×B - σE/ε
        
        参数:
            E: 电场张量
            B: 磁场张量
            J: 电流密度张量，可选
            dt: 时间步长
            dx, dy, dz: 空间步长
            
        返回:
            电场时间导数
        """
        c_squared = 1.0 / (self.eps0 * self.mu0)
        curl_B = self.compute_curl_B(B, dx, dy, dz)
        
        dE_dt = c_squared * curl_B
        
        # 添加电流项
        if J is not None:
            dE_dt = dE_dt - J / self.eps0
        
        # 添加电导率项
        if self.sigma > 0:
            dE_dt = dE_dt - self.sigma * E / self.eps0
        
        return dE_dt


class MHDEquations:
    """磁流体力学(MHD)方程实现"""
    
    def __init__(self, nu=1e-4, eta=1e-4, rho0=1.0):
        """初始化MHD方程参数
        
        参数:
            nu: 运动粘度
            eta: 磁扩散率
            rho0: 参考密度
        """
        self.nu = nu
        self.eta = eta
        self.rho0 = rho0
    
    def compute_laplacian(self, field, dx=1.0, dy=1.0, dz=1.0):
        """计算拉普拉斯算子 ∇²f
        
        参数:
            field: 场张量，形状为 [batch_size, x_dim, y_dim, z_dim, ...]
            dx, dy, dz: 空间步长
            
        返回:
            拉普拉斯算子结果，形状与输入相同
        """
        # 计算二阶偏导数
        d2f_dx2 = torch.gradient(torch.gradient(field, dim=1)[0], dim=1)[0] / (dx * dx)
        d2f_dy2 = torch.gradient(torch.gradient(field, dim=2)[0], dim=2)[0] / (dy * dy)
        d2f_dz2 = torch.gradient(torch.gradient(field, dim=3)[0], dim=3)[0] / (dz * dz)
        
        # 拉普拉斯算子
        laplacian = d2f_dx2 + d2f_dy2 + d2f_dz2
        
        return laplacian
    
    def compute_advection(self, velocity, field, dx=1.0, dy=1.0, dz=1.0):
        """计算对流项 (v·∇)f
        
        参数:
            velocity: 速度场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            field: 场张量，形状为 [batch_size, x_dim, y_dim, z_dim, ...]
            dx, dy, dz: 空间步长
            
        返回:
            对流项结果，形状与field相同
        """
        # 速度分量
        u = velocity[..., 0]
        v = velocity[..., 1]
        w = velocity[..., 2]
        
        # 场的梯度
        grad_f = torch.gradient(field, dim=(1, 2, 3))
        df_dx = grad_f[0] / dx
        df_dy = grad_f[1] / dy
        df_dz = grad_f[2] / dz
        
        # 对流项
        advection = u[..., None] * df_dx + v[..., None] * df_dy + w[..., None] * df_dz
        
        return advection
    
    def velocity_equation(self, velocity, magnetic_field, pressure=None, dx=1.0, dy=1.0, dz=1.0):
        """速度方程: ∂v/∂t + (v·∇)v = -∇p/ρ + (B·∇)B/ρ + ν∇²v
        
        参数:
            velocity: 速度场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            magnetic_field: 磁场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            pressure: 压力场，可选，形状为 [batch_size, x_dim, y_dim, z_dim]
            dx, dy, dz: 空间步长
            
        返回:
            速度时间导数
        """
        # 对流项
        advection = self.compute_advection(velocity, velocity, dx, dy, dz)
        
        # 粘性项
        viscous = self.nu * self.compute_laplacian(velocity, dx, dy, dz)
        
        # 磁场张力项 (B·∇)B
        magnetic_tension = self.compute_advection(magnetic_field, magnetic_field, dx, dy, dz)
        
        # 压力梯度项
        if pressure is not None:
            grad_p = torch.gradient(pressure, dim=(1, 2, 3))
            pressure_gradient = torch.stack([
                grad_p[0] / dx,
                grad_p[1] / dy,
                grad_p[2] / dz
            ], dim=-1) / self.rho0
        else:
            pressure_gradient = torch.zeros_like(velocity)
        
        # 速度方程
        dv_dt = -advection - pressure_gradient + magnetic_tension / self.rho0 + viscous
        
        return dv_dt
    
    def magnetic_equation(self, velocity, magnetic_field, dx=1.0, dy=1.0, dz=1.0):
        """磁场方程: ∂B/∂t = ∇×(v×B) + η∇²B
        
        参数:
            velocity: 速度场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            magnetic_field: 磁场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            磁场时间导数
        """
        # 计算v×B
        v_cross_B = torch.cross(velocity, magnetic_field, dim=-1)
        
        # 计算∇×(v×B)
        curl_v_cross_B = self.compute_curl(v_cross_B, dx, dy, dz)
        
        # 扩散项
        diffusion = self.eta * self.compute_laplacian(magnetic_field, dx, dy, dz)
        
        # 磁场方程
        dB_dt = curl_v_cross_B + diffusion
        
        return dB_dt
    
    def compute_curl(self, field, dx=1.0, dy=1.0, dz=1.0):
        """计算场的旋度 ∇×f
        
        参数:
            field: 场张量，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            场的旋度，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
        """
        # field = (fx, fy, fz)
        fx = field[..., 0]
        fy = field[..., 1]
        fz = field[..., 2]
        
        # 计算偏导数
        dfz_dy = torch.gradient(fz, dim=2)[0] / dy
        dfy_dz = torch.gradient(fy, dim=3)[0] / dz
        dfx_dz = torch.gradient(fx, dim=3)[0] / dz
        dfz_dx = torch.gradient(fz, dim=1)[0] / dx
        dfy_dx = torch.gradient(fy, dim=1)[0] / dx
        dfx_dy = torch.gradient(fx, dim=2)[0] / dy
        
        # 旋度分量
        curl_x = dfz_dy - dfy_dz
        curl_y = dfx_dz - dfz_dx
        curl_z = dfy_dx - dfx_dy
        
        return torch.stack([curl_x, curl_y, curl_z], dim=-1)
    
    def divergence_constraint(self, velocity, magnetic_field, dx=1.0, dy=1.0, dz=1.0):
        """计算散度约束条件
        
        参数:
            velocity: 速度场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            magnetic_field: 磁场，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            速度散度和磁场散度
        """
        # 速度散度 ∇·v
        div_v = self.compute_divergence(velocity, dx, dy, dz)
        
        # 磁场散度 ∇·B
        div_B = self.compute_divergence(magnetic_field, dx, dy, dz)
        
        return div_v, div_B
    
    def compute_divergence(self, field, dx=1.0, dy=1.0, dz=1.0):
        """计算场的散度 ∇·f
        
        参数:
            field: 场张量，形状为 [batch_size, x_dim, y_dim, z_dim, 3]
            dx, dy, dz: 空间步长
            
        返回:
            场的散度，形状为 [batch_size, x_dim, y_dim, z_dim]
        """
        # field = (fx, fy, fz)
        fx = field[..., 0]
        fy = field[..., 1]
        fz = field[..., 2]
        
        # 计算偏导数
        dfx_dx = torch.gradient(fx, dim=1)[0] / dx
        dfy_dy = torch.gradient(fy, dim=2)[0] / dy
        dfz_dz = torch.gradient(fz, dim=3)[0] / dz
        
        # 散度
        divergence = dfx_dx + dfy_dy + dfz_dz
        
        return divergence